import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.parameter import Parameter
from torch.nn.modules.module import Module


class MNN(nn.Module):
    def __init__(self, node_size, nhid0, nhid1, droput, alpha):
        # 声明神经网络都有哪些层
        # 这里的结构是：输入-encode0-encode1-decode0-decode1-输出

        # 于是如果我们输入的向量的长度为node_size，那么每一层神经元的个数分别是

        # 假设我们一次输入的batch_size = K
        # 输入 = [K, node_size]
        super(MNN, self).__init__()
        # encode0: [node_size, nhid0], 这一步输出 = [K, node_size] * [node_size, nhid0] = [K, nhid0]
        self.encode0 = nn.Linear(node_size, nhid0)
        # encode1: [nhid0, nhid1], 这一步输出 = [K, nhid0] * [nhid0, nhid1] = [K, nhid1]
        self.encode1 = nn.Linear(nhid0, nhid1)
        # decode0: [nhid1, nhid0], 这一步输出 = [K, nhid1] * [nhid1, nhid0] = [K, nhid0]
        self.decode0 = nn.Linear(nhid1, nhid0)
        # decode1: [nhid0, node_size], 这一步输出 = [K, nhid0] * [nhid0, node_size] = [K, node_size]
        # 自编码器希望最终的输出尽可能与输入相等
        self.decode1 = nn.Linear(nhid0, node_size)
        self.droput = droput
        # 论文公式(5)中的alpha
        self.alpha = alpha

    def forward(self, adj_batch, adj_mat, b_mat):
        # __init__中声明的神经网络的架构实现
        # 激活函数使用的是leaky_relu
        # 这里的结构是：输入-encode0-encode1-decode0-decode1-输出
        t0 = F.leaky_relu(self.encode0(adj_batch))
        t0 = F.leaky_relu(self.encode1(t0))
        # 这个中间结果使我们最终要保留的node embeeding
        embedding = t0
        t0 = F.leaky_relu(self.decode0(t0))
        t0 = F.leaky_relu(self.decode1(t0))
        embedding_norm = torch.sum(embedding * embedding, dim=1, keepdim=True)
        # 一阶相似度的loss值
        L_1st = torch.sum(adj_mat * (embedding_norm -
                                     2 * torch.mm(embedding, torch.transpose(embedding, dim0=0, dim1=1))
                                     + torch.transpose(embedding_norm, dim0=0, dim1=1)))
        # 二阶相似度的loss值
        # adj_batch是论文中的input X, t0是论文中最后输出的X^
        L_2nd = torch.sum(((adj_batch - t0) * b_mat) * ((adj_batch - t0) * b_mat))
        return L_1st, self.alpha * L_2nd, L_1st + self.alpha * L_2nd

    def savector(self, adj):
        t0 = self.encode0(adj)
        t0 = self.encode1(t0)
        return t0

# broadcast